#include "PhysicsWorld.h"

using namespace std;

void PhysicsWorld::processCollisions(dGeomID o1, dGeomID o2)
{
	dBodyID b1 = dGeomGetBody(o1);
	dBodyID b2 = dGeomGetBody(o2);

	if(b1 && b2 && dAreConnectedExcluding(b1, b2, dJointTypeContact)) 
		return;

	bool condFingerBoxContact;
	bool condBoxTableContact;
	bool condFingerTableContact;
	
	condFingerBoxContact = ((o1 == fingerGeom) && (o2 == magicBoxGeom)) || ((o2 == fingerGeom) && (o1 == magicBoxGeom));
	condBoxTableContact = ((o1 == tableGeom) && (o2 == magicBoxGeom)) || ((o2 == tableGeom) && (o1 == magicBoxGeom));
	condFingerTableContact = ((o1 == tableGeom) && (o2 == fingerGeom)) || ((o2 == tableGeom) && (o1 == fingerGeom));

    if(condFingerBoxContact)
	{
		int n = dCollide(o1, o2, 1, &this->contact.geom, sizeof(dContact));
		if(n>0)
		{
			for(int i = 0; i < n; i++)
			{
				this->contact.surface.mode = dContactSoftERP | dContactSoftCFM;
				this->contact.surface.mu = Macro::s_mu1;//dInfinity;
				this->contact.surface.soft_erp = 0.1;
				this->contact.surface.soft_cfm = 1e-4;
				this->conJoint = dJointCreateContact(world, contactgroup, &this->contact);
				dJointAttach(conJoint, b1, b2);
				
				if(dJointGetFeedback(conJoint) == NULL)
					dJointSetFeedback(conJoint, &fbCon);

				this->conInfo.status = 1;
			}
		}
	}

	else
	{
		static const int N = 4;
		dContact contact[N];
		int n = dCollide(o1, o2, N, &contact[0].geom, sizeof(dContact));
		if(n>0)
		{
			for(int i=0; i<n; i++)
			{
				contact[i].surface.mode = dContactSoftERP | dContactSoftCFM;
				contact[i].surface.mu = Macro::s_mu2;//dInfinity;
				contact[i].surface.soft_erp = 0.2;
				contact[i].surface.soft_cfm = 1e-4;
				dJointID c = dJointCreateContact(world, contactgroup, &contact[i]);
				dJointAttach(c, b1, b2);
			}
		}
	}
}
void nearCallback (void *data, dGeomID o1, dGeomID o2)
{
	unsigned long cat1 = dGeomGetCategoryBits (o1);
	unsigned long cat2 = dGeomGetCategoryBits (o2);
	unsigned long col1 = dGeomGetCollideBits (o1);
	unsigned long col2 = dGeomGetCollideBits (o2);
	
	if ((cat1 & col2) || (cat2 & col1))
		((PhysicsWorld*)data)->processCollisions(o1, o2);
}

PhysicsWorld::PhysicsWorld()
{
	dInitODE();

	if(!world)
		dWorldDestroy (world);
	if(!contactgroup)
		dJointGroupDestroy (contactgroup);
	if(!mousegroup)
		dJointGroupDestroy (mousegroup);
	if(!armgroup)
		dJointGroupDestroy (armgroup);
	if(!space)
		dSpaceDestroy (space);

	world = dWorldCreate();
	space = dHashSpaceCreate(0);
	contactgroup = dJointGroupCreate(0);
	mousegroup = dJointGroupCreate(0);
	armgroup = dJointGroupCreate(0);

	dWorldSetContactSurfaceLayer(world,0.001);
	dWorldSetContactMaxCorrectingVel(world, 0.1);

	dWorldSetGravity (this->world, 0, 0, GRAVITY);
	dWorldSetCFM (this->world, WCFM);
	dWorldSetERP (this->world, WERP);
	dWorldSetQuickStepNumIterations (world, 64);
	
	dSpaceSetCleanup(this->space, 1);

	groundPlane = dCreatePlane (this->space,0, 0, 1, 0);

	bodyNum = 0;
	geomNum = 0;

	mouseBody1 = NULL;
	mouseBody2 = NULL;
	mouseJoint = NULL;

	armJointNum = 0;

	timeStep = SIMSTEP;
}

PhysicsWorld::~PhysicsWorld()
{
	dJointGroupDestroy (this->contactgroup);
	dJointGroupDestroy (this->mousegroup);
	dJointGroupDestroy (this->armgroup);
	dSpaceDestroy (this->space);
	dWorldDestroy (this->world);
	dCloseODE();
}

void PhysicsWorld::servo(std::vector<Joint, Eigen::aligned_allocator<Joint>>& joints)
{
	int i = 0;
	for(std::vector<Joint, Eigen::aligned_allocator<Joint>>::iterator iter = joints.begin(); iter != joints.end(); iter++)
	{
		dJointID jointID = iter->id;
		dJointType jointType = iter->type;

		if(jointType == dJointTypeHinge)
		{
			AngleAxisd jointAA(iter->quat);

			dReal gain = Macro::s_gain;
			dReal maxForce = Macro::s_maxForce;

			dVector3 axis;
			dJointGetHingeAxis(jointID, axis);
			Vector3d trueAxis(axis);
			Vector3d desiredAxis(jointAA.axis());

			dReal trueAngle = dJointGetHingeAngle(jointID);
			dReal desiredAngle = jointAA.angle();

			if((trueAxis - (-desiredAxis)).norm() < POS_THRESH)
				desiredAngle *= -1;

			dReal error = trueAngle - desiredAngle;

			dReal desiredVelocity = -error * gain;

			if(Macro::s_sensorNoise && Macro::s_replay)
			{
				std::random_device rd;
				std::default_random_engine generator(rd());
				std::normal_distribution<double> distribution(0.0, 0.0007);
				while(1)
				{
					double gaussianNoise = distribution(generator);
					if(gaussianNoise < 0.0001 && gaussianNoise > -0.0001)
					{
						desiredVelocity += gaussianNoise;
						break;
					}
				}
			}

			dJointSetHingeParam(jointID, dParamFMax, maxForce);
			dJointSetHingeParam(jointID, dParamVel, desiredVelocity);
			dJointSetHingeParam(jointID, dParamFudgeFactor, 0.001);
			
			dBodySetFiniteRotationAxis(dJointGetBody(jointID, 0), axis[0], axis[1], axis[2]);
		}
	}
}

void PhysicsWorld::storePathPoint()
{
	dReal timeStep = Macro::s_timeStep;
	const dReal* pos = dBodyGetPosition(endEffector);
	const dReal* rot = dBodyGetQuaternion(endEffector);
	dVector3 mousePos = {Macro::s_mousePos[0], Macro::s_mousePos[1], Macro::s_mousePos[2]};
	PathPoint pathPoint(timeStep, pos, rot, mousePos);
	traj->pathPoints.push_back(pathPoint);
}

void PhysicsWorld::storeTotalInfo()
{
	dBodyID wristBody = dJointGetBody(wristJoint, 0);
	
	dReal timeStep = Macro::s_timeStep;
	dMass m;
	dBodyGetMass(magicBoxBody, &m);
	dReal weight = m.mass;
	dReal mu = Macro::s_mu2;
	Vector3d wristForce(fbWrist.f1);
	Vector3d wristTorque(fbWrist.t1);
	Vector3d wristPos(dBodyGetPosition(wristBody));
	Quaternion<double> wristRot(dBodyGetQuaternion(wristBody));
	int status = conInfo.status;
	Vector3d conForce = conInfo.force;

	TotalInfo totalInfo(timeStep, weight, mu, wristForce, wristTorque, wristPos, wristRot, status, conForce);
	traj->totalInfos.push_back(totalInfo);
}

void PhysicsWorld::storeContactInfo()
{
	//conInfo.reset();
	if(conInfo.status == 1)
	{
		dContactGeom geom = this->contact.geom;

		if(geom.g1 == fingerGeom)
			this->conInfo.force = Vector3d(fbCon.f1);
		else if(geom.g2 == fingerGeom)
			this->conInfo.force = Vector3d(fbCon.f2);
		this->conInfo.pos = Vector3d(geom.pos);
	}
}

void PhysicsWorld::animatePhysics(std::vector<Joint, Eigen::aligned_allocator<Joint>>& joints)
{
	conInfo.reset();
	int iter = 0.02 / timeStep;

	if(dJointGetFeedback(wristJoint) == NULL)
		dJointSetFeedback(wristJoint, &fbWrist);

	for(int i = 0; i < iter; i++)
	{
		servo(joints);

		dWorldQuickStep(this->world, timeStep);
		storeContactInfo();
		dJointGroupEmpty (this->contactgroup);
		dSpaceCollide(this->space, this, &nearCallback);

		if(Macro::s_mouseControl && Macro::s_record)
		{
			Macro::s_timeStep += timeStep;
			storePathPoint();
			storeTotalInfo();
		}

		if(Macro::s_replay && Macro::s_record)
		{
			Macro::s_timeStep += timeStep;
			storeTotalInfo();
		}
	}
}

void PhysicsWorld::reset()
{
	dJointGroupEmpty(armgroup);
	dJointGroupEmpty(mousegroup);

	for(std::map<dGeomID, dBodyID>::iterator iter = geoms.begin(); iter != geoms.end(); iter++)
	{
		dGeomID geomID = iter->first;
		dBodyID bodyID = iter->second;
		if(geomID)
			dGeomDestroy(geomID);
		if(bodyID)
			dBodyDestroy(bodyID);
	}

	bodies.clear();
	bodyNum = 0;

	geoms.clear();
	geomNum = 0;
	
	armJoints.clear();
	armJointNum = 0;
}